#include <stdlib.h>

#include <numeric>
#include <limits>
#include <string>
#include <sstream>
#include <random>

#ifndef _USE_MATH_DEFINES
#define _USE_MATH_DEFINES
#include <math.h>
#endif

#include "pso.h"
#include "pso_utils.h"

#include "CImg.h"
using namespace cimg_library;

extern std::vector<std::pair<double,double>>  linechart_in;
//extern std::vector<std::pair<double,double>>* linechart;

extern CImg<unsigned char> img_rgb, seg_face, img_grd;
extern CImg<unsigned char>* img_bg;
extern int Dx, Dy;

//----------------------------------------------------------------------------
// PSO parameters
//----------------------------------------------------------------------------
extern int Dim;
extern int numParticles;
extern int numIterations;
extern std::vector<double> BoundPosMin;
extern std::vector<double> BoundPosMax;
extern std::vector<double> BoundVeloMin;
extern std::vector<double> BoundVeloMax;
//----------------------------------------------------------------------------



//----------------------------------------------------------------------------
// PSO constrain by region
//----------------------------------------------------------------------------
extern int numConstrains;
extern double region_factor;
extern CImg<int>* region;
//----------------------------------------------------------------------------

//----------------------------------------------------------------------------
// PSO constrain by scales
//----------------------------------------------------------------------------
extern double user_scalex;  // 0.7% of width of image
extern double user_scaley;  // 0.5% of height of image
extern double user_tolx;  // tolerance of user preferred scale
extern double user_toly;  // tolerance of user preferred scale
//----------------------------------------------------------------------------


//----------------------------------------------------------------------------
// PSO constrain by region
//----------------------------------------------------------------------------
extern double edge_factor;
//----------------------------------------------------------------------------


//----------------------------------------------------------------------------
// PSO constrain by color
//----------------------------------------------------------------------------
extern double color_factor;
extern std::vector<int> clr_picked;
extern std::vector<int> loc_picked;
extern std::vector<std::vector<double>> gkernel;
extern int kernel_radius;
//----------------------------------------------------------------------------



//----------------------------------------------------------------------------
// PSO parameters
//----------------------------------------------------------------------------
// - All parameter values are set to values recommended b a research paper
//   that investigated the effects of various PSO parameter values on a set of
//   benchmark minimization problems.
// - Alternatively, the inertia weight can be dynamically vary according to 
//   the iterations.
// - If you set the cognitive value to be larger than the value of social weight
//   you place more weight on a particle's best known position that on the 
//   swarm's lobal best known position, and vice versa.
//----------------------------------------------------------------------------
extern double  w;// inertia weight
extern double c1;  // cognitive weight
extern double c2;  // social weight
// ----------------------------------------------------------------------------


//-----------------------------------------------------------------------------
// Region constrain
// ----------------------------------------------------------------------------
double is_selected(int x, int y) {
	if (x < 0.0 || x >= region->width() || y < 0.0 || y >= region->height())
		return 0.0;
	

	return (*region)(x,y);
}

double score_region(double x, double y)
{
	if (x < 0.0 || x >= region->width() || y < 0.0 || y >= region->height())
		return 0.0;

	int 
		lx = (int)floor(x),
		hx = (int)ceil(x),
		ly = (int)floor(y),
		hy = (int)ceil(y);

	double
		alpha = x - (double)lx,
		beta = y - (double)ly;

	double
		scorex0 = is_selected(lx,ly)*(1.0 - alpha) + is_selected(hx,ly)*alpha,
		scorex1 = is_selected(lx,hy)*(1.0 - alpha) + is_selected(hx,hy)*alpha,
		scorey0 = is_selected(lx,ly)*(1.0 - beta)  + is_selected(lx,hy)*beta,
		scorey1 = is_selected(hx,ly)*(1.0 - beta)  + is_selected(hx,hy)*beta;

	return (scorex0 + scorex1 + scorey0 + scorey1) / 4.0;
}

double constrain_by_region(const std::vector<std::pair<double,double>>& l)
{
	// NOTE: The default region is image boundary 
	// when user does not specify any preferred region.
	double score = 0.0;

	for (unsigned int i=0; i<l.size(); i++) {
		score += score_region(l[i].first,l[i].second);
	}

	return score/(double)l.size();
}
// ----------------------------------------------------------------------------

// ----------------------------------------------------------------------------
// Edge constrain
// ----------------------------------------------------------------------------
double is_edge(int x, int y) {
	return (double)(img_grd(x,y,0,0) > 0); 
}

double score_edge(double x, double y) {
	double score = 0.0;

	if (x < 0.0 || x >= img_grd.width() || 
		y < 0.0 || y >= img_grd.height())
		return 0.0;

	int 
		lx = (int)floor(x),
		hx = (int)ceil(x),
		ly = (int)floor(y),
		hy = (int)ceil(y);

	double
		alpha = x - (double)lx,
		beta = y - (double)ly;

	double
		scorex0 = is_edge(lx,ly)*(1.0 - alpha) + is_edge(hx,ly)*alpha,
		scorex1 = is_edge(lx,hy)*(1.0 - alpha) + is_edge(hx,hy)*alpha,
		scorey0 = is_edge(lx,ly)*(1.0 - beta)  + is_edge(lx,hy)*beta,
		scorey1 = is_edge(hx,ly)*(1.0 - beta)  + is_edge(hx,hy)*beta;

	return (scorex0 + scorex1 + scorey0 + scorey1) / 4.0;

	return score;
}

double constrain_by_edge(const std::vector<std::pair<double,double>>& l) {
	double score = 0.0;

	for (unsigned int i=0; i<l.size(); i++) {
		score += score_edge(l[i].first, l[i].second);
	}

	return score / (double)l.size();
}
// ----------------------------------------------------------------------------


// ----------------------------------------------------------------------------
// Color constrain
// ----------------------------------------------------------------------------
void get_avg_RGB(int x, int y, double& Red, double& Green, double& Blue){
	
	int R = kernel_radius;

	int x0 = x  - R;
	int y0 = y - R;
	int x1 = x  + R;
	int y1 = y + R;

	Red=0.0; Green=0.0; Blue=0.0;
	int gx=0, gy=0;
	for (int y=y0; y<=y1; y++, gy++) {
		for (int x=x0; x<=x1; x++, gx++) {
			if (x0>=0 && x1<Dx && y0>=0 && y1< Dy){
				Red   += (double)img_rgb(x,y,0,0) * gkernel[gy][gx];
				Green += (double)img_rgb(x,y,0,1) * gkernel[gy][gx];
				Blue  += (double)img_rgb(x,y,0,2) * gkernel[gy][gx];
			}
		}
		gx = 0;
	}

	Red = (Red < 0.0) ? 0.0: Red;
	Red = (Red >= 255.0) ? 255.0: Red;
	Green = (Green < 0.0) ? 0.0: Green;
	Green = (Green >= 255.0) ? 255.0: Green;
	Blue = (Blue < 0.0) ? 0.0: Blue;
	Blue = (Blue >= 255.0) ? 255.0: Blue;
}

double color_difference(int x, int y)
{
	double R,G,B;
	get_avg_RGB(x,y,R,G,B);

	double min_dist = DBL_MAX;
	for (unsigned int i=0; i<clr_picked.size()/3; i++) {
		double rR = (double)clr_picked[3*i + 0];
		double rG = (double)clr_picked[3*i + 1];
		double rB = (double)clr_picked[3*i + 2];

		double dR2 = (rR - R)*(rR - R);
		double dG2 = (rG - G)*(rG - G);
		double dB2 = (rB - B)*(rB - B);

		double dist = std::sqrt(dR2 + dG2 + dB2);
		if (min_dist > dist)
			min_dist = dist;
	}
	return std::exp(-min_dist/100.0);
}


double score_color(double x, double y) {
	double score = 0.0;

	if (x < 0.0 || x >= img_grd.width() || 
		y < 0.0 || y >= img_grd.height())
		return 0.0;

	int 
		lx = (int)floor(x),
		hx = (int)ceil(x),
		ly = (int)floor(y),
		hy = (int)ceil(y);

	double
		alpha = x - (double)lx,
		beta = y - (double)ly;

	double
		scorex0 = color_difference(lx,ly)*(1.0 - alpha) + color_difference(hx,ly)*alpha,
		scorex1 = color_difference(lx,hy)*(1.0 - alpha) + color_difference(hx,hy)*alpha,
		scorey0 = color_difference(lx,ly)*(1.0 - beta)  + color_difference(lx,hy)*beta,
		scorey1 = color_difference(hx,ly)*(1.0 - beta)  + color_difference(hx,hy)*beta;

	return (scorex0 + scorex1 + scorey0 + scorey1) / 4.0;
}

double constrain_by_color(const std::vector<std::pair<double,double>>& l) {
	double score = 0.0;

	for (unsigned int i=0; i<l.size(); i++) {
		score += score_color(l[i].first, l[i].second);
	}

	return score / (double)l.size();
}
// ----------------------------------------------------------------------------




void recon_linechart(
	std::vector<std::pair<double,double>>& l,
	const double& tx,
	const double& ty,
	const double& sx,
	const double& sy)
{
	l.resize(linechart_in.size());

	// scaling and then translate
	// - Assume that the input linechart always starts from the origin
	//   (the left-top corner of the image)
	for (unsigned int i=0; i<l.size(); i++) {
		l[i].first = (sx * linechart_in[i].first) + tx;
		l[i].second = (sy * linechart_in[i].second) + ty;
	}
}

static double get_scale_penalty(const double& sx, const double& sy)
{
	double 
		gauss_sx = (sx-user_scalex)*(sx-user_scalex) / (2.0*user_tolx*user_tolx),
		gauss_sy = (sy-user_scaley)*(sy-user_scaley) / (2.0*user_toly*user_toly);
	return 
		std::exp( -1.0 * (gauss_sx + gauss_sy) );
}

double PSO::ObjectiveFunc(const std::vector<double>& X)
{
	std::vector<std::pair<double,double>> l;
	recon_linechart(l,X[0],X[1],X[2],X[3]);

	// Compute score constrained by region user specified
	double s0 = constrain_by_region(l);

	// Compute score constrained by edge
	double s1 = constrain_by_edge(l);

	// Compute score constrained by color
	double s2 = constrain_by_color(l);

	// Total
	double sTotal = s0*region_factor + s1*edge_factor + s2*color_factor;

	// Scale panalty
	double sp = get_scale_penalty(X[2],X[3]);

	return sp * sTotal;
}

static void init_swarm(
	std::vector<PSO::Particle*>& swarm,
	double& bestGlobalFitness,
	std::vector<double>& bestGlobalPosition,
	Utils::Random& ran)
{
	for (int p=0; p<numParticles; p++) {
		std::vector<double> ranPos(Dim);
		for (int i=0; i<Dim; i++) {
			double lo = BoundPosMin[i];
			double hi = BoundPosMax[i];
			ranPos[i] = (hi - lo) * ran.NextDouble() + lo;
		}

		double fitness = PSO::ObjectiveFunc(ranPos);

		std::vector<double> ranVelo(Dim);
		for (int i=0; i<Dim; i++) {
			double lo = BoundVeloMin[i];
			double hi = BoundVeloMax[i];
			ranVelo[i] = (hi - lo) * ran.NextDouble() + lo;
		}

		swarm[p] = new PSO::Particle(
			ranPos,fitness,ranVelo,ranPos,fitness);

		if (swarm[p]->get_fitness() > bestGlobalFitness) {
			bestGlobalFitness = swarm[p]->get_fitness();
			swarm[p]->copy_position_to(bestGlobalPosition);
		}
	}
}


void app_pso(bool show)
{
	Utils::Random ran;

	CImg<unsigned char> canvas;
	CImgDisplay disp;
	const unsigned char 
		red[] = {255,0,0}, blue[] = {0,0,255}, 
		black[] = {0,0,0}, white[] = {255,255,255};

	if (show) {
		canvas.assign(*img_bg);
		cimg_forXYC(canvas,x,y,c) {
			if ((*region)(x,y)==1) continue;
			canvas(x,y,c) = canvas(x,y,c)*0.7f + seg_face(x,y,c)*0.3f;
		}
		canvas.display(disp);
		disp.set_title("Particle Swarm Optimization");
	}

	std::vector<PSO::Particle*> swarm;
	std::vector<double> bestGlobalPosition;
	double bestGlobalFitness = DBL_MIN;

	swarm.resize(numParticles);
	bestGlobalPosition.resize(Dim);

	init_swarm(swarm,bestGlobalFitness,bestGlobalPosition,ran);

	if (show) {
		draw_swarm(swarm, canvas, &bestGlobalPosition);
		canvas.display(disp);
	}

	std::vector<double> newVelo;
	std::vector<double> newPos;
	double              newFitness;

	newVelo.resize(Dim);
	newPos.resize(Dim);

	// Main processing loop
	for (int iter = 0; iter < numIterations; iter++) {
		for (int p=0; p<numParticles; p++) {
			PSO::Particle* currP = swarm[p];

			for (int j=0; j<Dim; j++) {
				double r1 = ran.NextDouble();
				double r2 = ran.NextDouble();

				double 
					local = currP->get_best_position(j) - currP->get_position(j),
					global = bestGlobalPosition[j] - currP->get_position(j);

				newVelo[j] = 
					(w*currP->get_velocity(j)) + (c1*r1*local) + (c2*r2*global);

				if (newVelo[j] < BoundVeloMin[j])
					newVelo[j] = BoundVeloMin[j];
				else if (newVelo[j] > BoundVeloMax[j])
					newVelo[j] = BoundVeloMax[j];
			}// Each j

			currP->copy_velocity_from(newVelo);

			for (int j=0; j<Dim; j++) {
				newPos[j] = currP->get_position(j) + newVelo[j];
				if (newPos[j] < BoundPosMin[j])
					newPos[j] = BoundPosMin[j];
				else if (newPos[j] > BoundPosMax[j])
					newPos[j] = BoundPosMax[j];
			}

			currP->copy_position_from(newPos);

			newFitness = PSO::ObjectiveFunc(newPos);
			currP->set_fitness(newFitness);

			if (newFitness > currP->get_best_fitness()) {

				currP->copy_best_position_from(newPos);
				currP->set_best_fitness(newFitness);
			}

			if (newFitness > bestGlobalFitness) {
				for (int i=0; i<Dim; i++)
					bestGlobalPosition[i] = newPos[i];
				bestGlobalFitness = newFitness;
			}
		} // Each Particle

		if (show) {
			canvas.assign(*img_bg);
			cimg_forXYC(canvas,x,y,c) {
				if ((*region)(x,y)==1) continue;
				canvas(x,y,c) = canvas(x,y,c)*0.7f + seg_face(x,y,c)*0.3f;
			}
			draw_swarm(swarm, canvas, &bestGlobalPosition);
			draw_linechart(bestGlobalPosition, linechart_in, canvas);

			std::stringstream ss;
			ss << iter << " / " << numIterations << "\n"
			   << "Best Fit: " << bestGlobalFitness << "\n"
			   << "Best Sol: " << bestGlobalPosition[0] << " , " << bestGlobalPosition[1] << " , "
			                   << bestGlobalPosition[2] << " , " << bestGlobalPosition[3];
			canvas.draw_text(20,20,ss.str().c_str(),white,black,0.7f,15);

			canvas.display(disp);
			disp.wait(100);
		}

	} // Main PSO loop

	for (unsigned int i=0; i<swarm.size(); i++)
		delete swarm[i];

	getchar();
}


























